# 方法一 双O（mn）
class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        m, n = len(matrix), len(matrix[0])
        dp = [[0] * n for _ in range(m)]
        ans = 0
        for i in range(m):
            for j in range(n):
                if matrix[i][j] == 1:
                    # 动态规划
                    dp[i][j] = min(dp[i][j - 1], dp[i - 1][j], dp[i - 1][j - 1]) + 1
                    ans += dp[i][j]
        '''例dp数组
        [[0, 1, 1, 1], 
        [1, 1, 2, 2], 
        [0, 1, 2, 3]]
        '''
        return ans

# 方法二 优化 空间为O(1)
class Solution:
    def countSquares(self, matrix: List[List[int]]) -> int:
        m, n = len(matrix), len(matrix[0])
        ans = 0
        for i in range(m):
            for j in range(n):
                if matrix[i][j] == 1:
                    # 处理边界情况
                    if i - 1 < 0 or j - 1 < 0:
                        matrix[i][j] = 1
                    else:
                        # 动态规划
                        matrix[i][j] = min(matrix[i][j - 1], matrix[i - 1][j], matrix[i - 1][j - 1]) + 1
                    ans += matrix[i][j]
        return ans